#!/usr/bin/env python
# coding: utf-8

# In[3]:


pip install tushare

# In[6]:


import tushare as ts
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score

# 1. 获取股票数据
def get_stock_data(ts_code, start_date, end_date, token):
    ts.set_token(token)
    pro = ts.pro_api()
    df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
    df = df.set_index('trade_date')
    df.index = pd.to_datetime(df.index)
    df = df[['open', 'high', 'low', 'close', 'vol']]
    df.rename(columns={'vol': 'volume'}, inplace=True)
    return df

# 2. 计算技术指标
def calculate_technical_indicators(df):
    df['MA5'] = df['close'].rolling(window=5).mean()
    df['MA10'] = df['close'].rolling(window=10).mean()
    df['RSI'] = df['close'].rolling(window=14).apply(lambda x: x.diff().mean() / x.diff().abs().mean() * 100)
    df['MOM'] = df['close'].diff(periods=5)
    df['EMA12'] = df['close'].ewm(span=12, adjust=False).mean()
    df['EMA26'] = df['close'].ewm(span=26, adjust=False).mean()
    df['MACD'] = df['EMA12'] - df['EMA26']
    df['MACDsignal'] = df['MACD'].ewm(span=9, adjust=False).mean()
    df['MACDhist'] = df['MACD'] - df['MACDsignal']
    df.dropna(inplace=True)
    return df

# 3. 准备数据
def prepare_data(df):
    df['price_change'] = df['close'].diff(periods=1)
    df['target'] = np.where(df['price_change'].shift(-1) > 0, 1, -1)
    features = ['close', 'volume', 'MA5', 'MA10', 'RSI', 'MOM', 'EMA12', 'EMA26', 'MACD', 'MACDsignal', 'MACDhist']
    X = df[features]
    y = df['target']
    split = int(len(df) * 0.9)
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = y[:split], y[split:]
    return X_train, X_test, y_train, y_test

# 4. 建立随机森林模型
def build_model(X_train, y_train):
    model = RandomForestClassifier(max_depth=3, n_estimators=10, min_samples_leaf=10, random_state=123)
    model.fit(X_train, y_train)
    return model

# 5. 模型评估
def evaluate_model(model, X_test, y_test):
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(f'Model Accuracy: {accuracy:.2f}')
    return y_pred

# 6. 计算收益率
def calculate_return(df, y_pred, split):
    df['prediction'] = np.nan  # 初始化 prediction 列为 NaN
    df.iloc[split:, df.columns.get_loc('prediction')] = y_pred  # 只在测试集部分填充预测值
    df['p_change'] = df['close'].pct_change()
    df['origin'] = (1 + df['p_change']).cumprod()
    df['strategy'] = (1 + df['prediction'].shift(1) * df['p_change']).cumprod()
    return df

# 主函数
def main():
    ts_code = '000001.SZ'  # 示例股票：平安银行
    start_date = '20200101'
    end_date = '20250101'
    token = '1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c'  # 替换为你的 Tushare Token

    # 获取数据
    df = get_stock_data(ts_code, start_date, end_date, token)
    
    # 计算技术指标
    df = calculate_technical_indicators(df)
    
    # 准备数据
    X_train, X_test, y_train, y_test = prepare_data(df)
    
    # 建立模型
    model = build_model(X_train, y_train)
    
    # 模型评估
    y_pred = evaluate_model(model, X_test, y_test)
    
    # 计算收益率
    df = calculate_return(df, y_pred, split=int(len(df) * 0.9))
    
    # 打印收益率结果
    print("Original Strategy Cumulative Return: ", df['origin'].iloc[-1])
    print("Model Strategy Cumulative Return: ", df['strategy'].iloc[-1])

if __name__ == "__main__":
    main()






